from ..data_processing.MD17 import MD17
from torch_geometric.data import DataLoader
import torch
 
dataset = MD17(root='dataset/', name='aspirin') #ethanol aspirin benzene_old
split_idx = dataset.get_idx_split(len(dataset.data.y), train_size=2000, valid_size=100, seed=42)
print(split_idx)
train_dataset, valid_dataset, test_dataset = dataset[split_idx['train']], dataset[split_idx['valid']], dataset[split_idx['test']]
train_loader = DataLoader(train_dataset, batch_size=10,shuffle=False)
generator_data = iter(train_loader)

atom=21
index=50
diff=[]
s2=[]
diff1=[]
s21=[]
while index>0:
    data = next(generator_data)
    # print(data.y/23.06052)
    i=1
    x1=data.pos[atom*(i-1):atom*i]
    f1=data.force[atom*(i-1):atom*i]
    i=2
    x2=data.pos[atom*(i-1):atom*i]
    f2=data.force[atom*(i-1):atom*i]
    i=3
    x3=data.pos[atom*(i-1):atom*i]
    f3=data.force[atom*(i-1):atom*i]


    x2=x2+torch.mean(x1-x2)
    x3=x3+torch.mean(x1-x3)
 
    sigma_2= torch.sum((f1-f2)*(x2-x1))/torch.sum(torch.square(f1-f2))
    print("1:",sigma_2)
    x0_1=(sigma_2*(f1+f2)+x1+x2)/2

    sigma_2= torch.sum((f2-f3)*(x3-x2))/torch.sum(torch.square(f2-f3))
    x0_2=(sigma_2*(f2+f3)+x2+x3)/2
    print("2:",sigma_2)

    diff.append(torch.sum(x3-x2))
    s2.append(torch.sqrt(torch.square(x3-x2).sum()).item())
    diff1.append(torch.sum(x0_2-x0_1))
    s21.append(torch.sqrt(torch.square(x0_2-x0_1).sum()).item())

    index=index-1;

# print(diff[0:10])
# print(diff1[0:10])
print(torch.tensor(s2).mean())
print(torch.tensor(s21).mean())
print(torch.tensor(0)==0)
